# -*- coding: utf-8 -*-
"""
Created on Wed Sep  4 16:06:52 2024

@author: jzz
"""

import os
import h5py
from matplotlib import pyplot as plt
from tqdm import tqdm
import copy
import database_new_student as database_new
import re
import pandas as pd
import numpy as np
import sys
import datetime
import warnings
sys.path.append('D:\oneDrive\self_use_program')


yunwind = database_new.YunWind()  

def get_index_members(concept_index):
    concept_index_sector_dict = {}  # 字典中存储每个index的成分股
    for i in range(concept_index.shape[0]):
        row=concept_index.iloc[i,:]
        index_code = row['S_INFO_WINDCODE']
        concept_stock = yunwind.read_sql("select * from AINDEXMEMBERSWIND where F_INFO_WINDCODE='{}'".format(index_code))
        date_min = row['S_INFO_LISTDATE']  # 从发布日期开始
        today = datetime.datetime.today()
        stock_codes = sorted(list(concept_stock['S_CON_WINDCODE'].unique()))  # 提取所有出现过的成分股，删除重复值
        ind_res = pd.DataFrame(columns=stock_codes, index=pd.date_range(date_min, today))  # 创建一个index是从发布日期到现在
        for stock_code in stock_codes:  # 遍历每一个成分股（有的成分股可能会出现多次）
            stock_rows = concept_stock[concept_stock['S_CON_WINDCODE'] == stock_code]
            for j in range(stock_rows.shape[0]): 
                row1=stock_rows.iloc[j,:]
                if not pd.isnull(row1['S_CON_OUTDATE']):
                    loc = (ind_res.index >= row1['S_CON_INDATE']) & (ind_res.index < row1['S_CON_OUTDATE'])
                    ind_res.loc[loc, stock_code] = 1
                else:
                    loc = ind_res.index >= row1['S_CON_INDATE']
                    ind_res.loc[loc, stock_code] = 1
        concept_index_sector_dict[index_code] = ind_res
        print(index_code)
    return concept_index_sector_dict


# 取出各概念指数，并计算各指数的
concept_index = yunwind.read_sql("select * from AINDEXDESCRIPTION where S_INFO_INDEXCODE='647090000'") 


list(concept_index['S_INFO_INDEXSTYLE'].unique())
# 排除券商金股、中概股等指数 
concept_index = concept_index[concept_index['S_INFO_INDEXSTYLE'].isin(['题材', 'Wind热门概念指数', 'Wind热门概念指数|题材', 
                                                          '定制'])]
concept_index = concept_index[~concept_index['S_INFO_NAME'].str.contains('退市')]
concept_index = concept_index[~concept_index['S_INFO_NAME'].str.contains('首板|打板|连板')] 

concept_index = concept_index.drop(concept_index[concept_index['S_INFO_NAME'] == '近端次新股指数'].index)

concept_index = concept_index.drop(concept_index[concept_index['S_INFO_NAME'] == 'ST板块指数'].index)
'''
concept_index_sector_dict=get_index_members(concept_index)
# 处理字典的键
updated_dict = {f'_{key.replace(".", "_")}': df for key, df in concept_index_sector_dict.items()}
key_to_remove = '_884245_WI'
if key_to_remove in updated_dict:
    del updated_dict[key_to_remove]
'''

# 定义 HDF5 文件名
hdf5_filename = '各指数成分股.h5'

'''
# 使用 HDFStore 进行存储
with pd.HDFStore(hdf5_filename, mode='w') as store:
    for key, df in updated_dict.items():
        df=df.fillna(0)
        print(f"Storing DataFrame with key '{key}'")
        store.put(key, df, format='table')  # 使用 'table' 格式以支持查询
        # 你也可以使用 store[key] = df 直接存储，默认使用 'fixed' 格式
'''
with pd.HDFStore(hdf5_filename, mode='r') as store:
    # 获取文件中的所有键
    keys = store.keys()
    
    # 初始化一个空字典来存储 DataFrame
    index_members_dict = {}
    
    # 遍历所有键，读取对应的 DataFrame，并存储到字典中
    for key in keys:
        # key 是以 '/' 开头的路径，我们需要去掉开头的 '/'
        clean_key = key.lstrip('/')
        index_members_dict[clean_key] = store[key]
    
    
def get_new_in_sql(new_id,key='OBJECT_ID',table='ASHAREEODPRICES',cols=['*']):
    '''用join替代in函数'''
    forms='SELECT %s AS' + f' {key} ' + ' '.join(['UNION ALL SELECT %s']*(len(new_id)-1)) 
    sql0 = forms%tuple(new_id)
    
    cols = ['t1.'+col for col in cols]
    cols = ','.join(cols)
    
    sql = f'''
    SELECT {cols} FROM {table} t1 INNER JOIN ({sql0}) t2 ON t1.{key} = t2.{key};
            '''

    # sql = sql.format(table,sql0) 

    return sql 

# 取出指数行情 
new_id = ["'"+x+"'" for x in concept_index['S_INFO_WINDCODE']] 
sql = get_new_in_sql(new_id,key='S_INFO_WINDCODE',table='AIndexWindIndustriesEOD'.upper(),
                                 cols=['S_INFO_WINDCODE','TRADE_DT','S_DQ_OPEN','S_DQ_CLOSE']) 
index_data = yunwind.read_sql(sql) 

index_close = pd.pivot_table(index_data,index='TRADE_DT',columns='S_INFO_WINDCODE',values='S_DQ_CLOSE') 
index_close.index = pd.to_datetime(index_close.index) 


index_pct = index_close.pct_change() 

for index_code in index_pct.columns:
    row = concept_index[concept_index['S_INFO_WINDCODE'] == index_code]
    data = pd.to_datetime(row['S_INFO_LISTDATE'].item())
    index_pct[index_code] = np.where(index_pct.index < data, np.nan, index_pct[index_code])
    print(index_code,data)

    

def compare_indices_members(index1_code, index2_code, date):
    #改一下index_code的形式，因为dict中存储的key名不能有.
    index1_code = '_' + index1_code.replace('.', '_')
    index2_code = '_' + index2_code.replace('.', '_')
    # code1中的成分股
    index1_members = index_members_dict[index1_code]
    index1_members_date = index1_members.loc[date]
    result1 = index1_members_date[index1_members_date.eq(1)].index.tolist()
    
    # code2中的成分股
    index2_members = index_members_dict[index2_code]
    index2_members_date = index2_members.loc[date]
    result2 = index2_members_date[index2_members_date.eq(1)].index.tolist()
    
    # 相同元素个数
    common_elements = set(result1) & set(result2)
    common_count = len(common_elements)
    
    # 相同元素占code1、code2的比例
    code1_count = len(set(result1))
    code2_count = len(set(result2))
    
    proportion1 = common_count / code1_count if code1_count > 0 else 0
    proportion2 = common_count / code2_count if code2_count > 0 else 0
    
    return proportion1, proportion2



def find_sorted_concept_index(index_pct):
    '''对各概念指数的变化率进行排序，返回排序后的列名'''
    # 创建一个新的DataFrame用于存储每行排序后的列名
    sorted_indices = index_pct.apply(lambda row: row.sort_values(ascending=False).index.tolist(), axis=1)
    return sorted_indices

sorted_indices = find_sorted_concept_index(index_pct)
sorted_indices = sorted_indices.loc['2015-12-31':'2024-9-18']


def determine_best_index(sorted_indices, index_num_select):
    '''目前只写了选2个index的，还没写选3个及以上的。。。'''
    best_indices=pd.DataFrame(index=sorted_indices.index, columns=range(index_num_select))
    for i in range(sorted_indices.shape[0]):
        sorted_indices_row = sorted_indices.iloc[i]
        date = sorted_indices.index[i]
        for index_code in sorted_indices_row[1:]:#自己和自己不用比较
            index1_code = sorted_indices_row[0]
            #求相同元素占比
            proportion1 = compare_indices_members(index1_code, index_code, date)[0]
            proportion2 = compare_indices_members(index1_code, index_code, date)[1]
            if proportion1 > 0.5 or proportion2 > 0.5:
                print(i, index1_code, index_code)
                continue
            else:
                index2_code = index_code
                break
        best_indices.iloc[i, 0] = index1_code
        best_indices.iloc[i, 1] = index2_code
    return best_indices
    


# 找出当天最为强势的几个板块
index_num_select = 2
best_concept_index = determine_best_index(sorted_indices, index_num_select)


#把wind代码换成名字
def get_index_name(index_code):
    try:
        name = concept_index.loc[concept_index['S_INFO_WINDCODE'] == index_code, 'S_INFO_NAME'].values[0]
    except IndexError:
        name = np.nan  # 返回 NaN
    return name

best_concept_index_name = best_concept_index.copy()

for i in range(best_concept_index.shape[0]):  # 遍历行
    for j in range(best_concept_index.shape[1]):  # 遍历列
        code = best_concept_index.iloc[i, j]
        name = get_index_name(code)
        if not pd.isna(name):
            best_concept_index_name.iloc[i,j] = name
        print(name,i,j)

columns1 = []
for i in range(index_num_select):  # 每个 DataFrame 有 2 列
    columns1.append(best_concept_index.iloc[:, i])  # A 的第 i 列
    columns1.append(best_concept_index_name.iloc[:, i])  # B 的第 i 列

# 拼接这些列
result1 = pd.concat(columns1, axis=1)

result1.to_excel('最强势指数.xlsx')

stock_num_select = 2
'''
# 定义每批提取的行数
batch_size = 1000000
offset = 0

# 定义 SQL 查询的基础部分
base_sql = "SELECT S_INFO_WINDCODE, TRADE_DT, S_DQ_OPEN, S_DQ_CLOSE FROM ASHAREEODPRICES"

# 创建一个空的 DataFrame 用于存储结果
all_data = pd.DataFrame()

# 使用 tqdm 进度条
with tqdm(desc="Reading Data", unit="batch") as pbar:
    while True:
        # 构建完整的 SQL 查询，添加 LIMIT 和 OFFSET
        sql = f"{base_sql} LIMIT {batch_size} OFFSET {offset}"
        
        # 从数据库读取数据
        chunk = yunwind.read_sql(sql)
        
        # 如果没有读取到数据，退出循环
        if chunk.empty:
            break
        
        # 将读取的数据追加到总 DataFrame 中
        all_data = pd.concat([all_data, chunk], ignore_index=True)

        # 更新偏移量
        offset += batch_size
        
        # 更新进度条
        pbar.update(1)


grouped = all_data.groupby('S_INFO_WINDCODE')

# 创建一个字典，键为 S_INFO_WINDCODE 的值，值为对应的 DataFrame
df_dict = {windcode: group for windcode, group in grouped}
# 处理字典的键
df_updated_dict = {f'_{key.replace(".", "_")}': df for key, df in df_dict.items()}
key_to_remove = '_884245_WI'
if key_to_remove in df_updated_dict:
    del df_updated_dict[key_to_remove]
'''
# 定义 HDF5 文件名
hdf5_filename1 = 'A股数据.h5'

'''
# 使用 HDFStore 进行存储
with pd.HDFStore(hdf5_filename1, mode='w') as store:
    for key, df in df_updated_dict.items():
        df=df.fillna(0)
        print(f"Storing DataFrame with key '{key}'")
        store.put(key, df, format='table')  # 使用 'table' 格式以支持查询
        # 你也可以使用 store[key] = df 直接存储，默认使用 'fixed' 格式
'''

with pd.HDFStore(hdf5_filename1, mode='r') as store:
    # 获取文件中的所有键
    keys = store.keys()
    
    # 初始化一个空字典来存储 DataFrame
    all_stock_dict = {}
    
    # 遍历所有键，读取对应的 DataFrame，并存储到字典中
    for key in keys:
        # key 是以 '/' 开头的路径，我们需要去掉开头的 '/'
        clean_key = key.lstrip('/')
        all_stock_dict[clean_key] = store[key]
        print(key)



def get_stock_date_pct(stock_code, date, all_stock_dict):
    '''注意这里需要'''
    stock_code1 = '_' + stock_code.replace('.', '_')
    # 检查 stock_code1 是否在字典中，不考虑查不到的股票
    if stock_code1 not in all_stock_dict:
        return np.nan
    stock = all_stock_dict[stock_code1].copy()
    stock['TRADE_DT'] = pd.to_datetime(stock['TRADE_DT'])  # 转换为日期类型
    stock.set_index('TRADE_DT', inplace=True)
    stock = stock.sort_index()
    stock_pct = stock['S_DQ_CLOSE'].pct_change()
    try:
        return stock_pct.loc[date].item()
    except KeyError:  # 如果日期不存在
        return np.nan
    except ValueError:  # 如果索引存在但数据为空
        return np.nan

def find_best_concept_stock(date, key, index_members_dict, all_stock_dict, stock_num_select):
    # 获取对应的 DataFrame
    index_members = index_members_dict[key]  # 确保 df_dict 已正确定义
    date_row = index_members.loc[date]
    date_row1 = date_row.astype(float)
    
    for stock_code in date_row.index:
        if date_row[stock_code] == 1:
            pct_stock = get_stock_date_pct(stock_code, date, all_stock_dict)  
            date_row1[stock_code] = pct_stock
            print(pct_stock, stock_code)
        else:
            date_row1[stock_code] = np.nan

    # 选择最大的 stock_num_select 个值的列名
    top_columns = date_row1.nlargest(stock_num_select).index.tolist()
    
    return top_columns
find_best_concept_stock('2016-12-31', '_884025_WI', index_members_dict, all_stock_dict, stock_num_select)

def get_best_stock_columns(max_num):
    all_best_stocks=pd.DataFrame(index=best_concept_index.index, columns=range(stock_num_select))
    for index, value in best_concept_index.iloc[:, max_num].items():
        index_code = '_' + value.replace('.', '_')
        print(index,index_code)
        best_stocks = find_best_concept_stock(index, index_code, index_members_dict, all_stock_dict, stock_num_select)
        all_best_stocks.loc[index, :] = best_stocks
    return all_best_stocks


all_best_stocks=pd.DataFrame(index=best_concept_index.index, columns=range(stock_num_select*index_num_select))
for max_num in range(index_num_select):
    best_stock=get_best_stock_columns(max_num)
    start_col = max_num * stock_num_select
    end_col = start_col + stock_num_select
    # 将 best_stock 按列赋值到 all_best_stocks 的对应列
    all_best_stocks.iloc[:, start_col:end_col] = best_stock
  
all_stock_name = yunwind.read_sql("select S_INFO_WINDCODE, S_INFO_NAME from ASHAREDESCRIPTION")

def get_stock_name(stock_code):
    try:
        name = all_stock_name.loc[all_stock_name['S_INFO_WINDCODE'] == stock_code, 'S_INFO_NAME'].values[0]
    except IndexError:
        name = np.nan  # 返回 NaN
    return name

all_best_stocks = all_best_stocks.loc['2015-12-31':'2024-9-6']


        
        
def buy_sell_strategy(date, stock_code):
    stock_code1 = '_' + stock_code.replace('.', '_')
    stock = all_stock_dict[stock_code1].copy()
    stock['TRADE_DT'] = pd.to_datetime(stock['TRADE_DT'])  # 转换为日期类型
    stock.set_index('TRADE_DT', inplace=True)
    stock = stock.sort_index()
    buy_price = stock['S_DQ_CLOSE'].loc[date]#取当天的收盘价
    sell_price = stock['S_DQ_OPEN'].shift(-1).loc[date]#取后一天的开盘价
    gainrate = (sell_price - buy_price)/buy_price
    return gainrate, buy_price, sell_price
    
#删掉法定节假日
for index, row in all_best_stocks.iterrows():
    for j in range(len(row)):  # 遍历每行的每个值
        stock_code = row.iloc[j]
        try:
            gainrate = buy_sell_strategy(index, stock_code)
        except KeyError:
            all_best_stocks.drop(index, inplace=True)  # 删除对应的行
            break  # 退出当前行的循环，进入下一行

#把wind代码换成名字
all_best_stocks_name = all_best_stocks.copy()

    
for i in range(all_best_stocks.shape[0]):  # 遍历行
    for j in range(all_best_stocks.shape[1]):  # 遍历列
        code = all_best_stocks.iloc[i, j]
        name = get_stock_name(code)
        if not pd.isna(name):
            all_best_stocks_name.iloc[i,j] = name
        print(name,i,j)
        
all_best_stocks_gainrate = all_best_stocks.copy()
all_best_stocks_close = all_best_stocks.copy()
all_best_stocks_next_open = all_best_stocks.copy()

#记录收益
for i in range(all_best_stocks.shape[0]):  # 遍历行
    for j in range(all_best_stocks.shape[1]):  # 遍历列
        code = all_best_stocks.iloc[i, j]
        index = all_best_stocks.index[i]
        gainrate, close, next_open = buy_sell_strategy(index, code)
        all_best_stocks_gainrate.iloc[i,j] = gainrate
        all_best_stocks_close.iloc[i,j] = close
        all_best_stocks_next_open.iloc[i,j] = next_open
        print(gainrate,i,j)
        
        



columns = []
for i in range(stock_num_select*index_num_select):  # 每个 DataFrame 有 4 列
    columns.append(all_best_stocks.iloc[:, i])  
    columns.append(all_best_stocks_name.iloc[:, i])  
    columns.append(all_best_stocks_close.iloc[:, i])  
    columns.append(all_best_stocks_next_open.iloc[:, i])  
    columns.append(all_best_stocks_gainrate.iloc[:, i])  

# 拼接这些列
result = pd.concat(columns, axis=1)

# 重命名列名（可选）
result.columns = ['代码1','名称1','收盘价1', '后一天开盘价1', '收益率1','代码2','名称2','收盘价2', '后一天开盘价2','收益率2','代码3','名称3','收盘价3', '后一天开盘价3','收益率3','代码4','名称4','收盘价4', '后一天开盘价4','收益率4']
result.to_excel('最强势股票.xlsx')
result = result[:-1]

class Analyze:
    '''计算指标或画图进行分析''' 

    def evaluate_indicator(net_value,tradeday_per_year,anual_riskfree_rate):
        '''
        已知净值曲线之后进行各类指标的计算
    
        Parameters
        ----------
        net_value : pd.Series 
            时间序列.
        anual_riskfree_rate : float
            年化无风险利率.
        tradeday_per_year：如果净值序列是月度的，那就取12；如果是日度的，就取250（每年的交易日数据）；如果是周度的，就取52（每年的周数）
    
        Returns
        -------
        annual_yield：年化收益率
        annual_std：每期收益率的年化标准差
        sharpe_ratio：年化收益率的夏普比率
        drawdown：回撤曲线
        max_drawdown：最大回撤
    
        '''
        net_value0 = copy.copy(net_value)
        
        net_value = np.array(net_value)
        len_nav = len(net_value)
        annual_yield = (net_value[-1]/net_value[0])**(tradeday_per_year/len_nav)-1
        
        yield_of_nav = net_value[1:]/net_value[0:-1]-1
        annual_std = np.std(yield_of_nav)*np.sqrt(tradeday_per_year)
        
        sharpe_ratio = (annual_yield-anual_riskfree_rate)/annual_std
        
        max_nv = net_value0.cummax() 
        drawdown = (net_value0 - max_nv) / max_nv
        max_drawdown = np.max(-drawdown)
        
        # 分年度收益
        yield_year_last = net_value0.resample('y').last().pct_change() 
        yield_year_first_to_last = net_value0.resample('y').last()/net_value0.resample('y').first()-1
        
        return annual_yield,annual_std,sharpe_ratio,drawdown,max_drawdown,yield_year_last,yield_year_first_to_last 
    
    def evaluate_indicator_multi_asset(net_value,tradeday_per_year,anual_riskfree_rate):
        '''
        和evaluate_indicator区别是这里net_value是多资产的净值序列组成的dataframe

        Parameters
        ----------
        net_value : DataFrame
            DESCRIPTION.
        tradeday_per_year：如果净值序列是月度的，那就取12；如果是日度的，就取250（每年的交易日数据）；如果是周度的，就取52（每年的周数）
        anual_riskfree_rate : float
            0 or 0.03 ...

        Returns
        -------
        indicators : DataFrame
            DESCRIPTION.
        drawdowns : DataFrame
            DESCRIPTION.
        yield_year_lasts : DataFrame
            DESCRIPTION.
        yield_year_first_to_lasts : DataFrame
            DESCRIPTION.

        ''' 
        indicators = pd.DataFrame(index=net_value.columns,columns=['年化收益率','年化波动率','夏普比率','最大回撤'])
        drawdowns = [] 
        yield_year_lasts = [] 
        yield_year_first_to_lasts = [] 
        
        for code in net_value.columns:
            cur_nv = net_value[code]
            annual_yield,annual_std,sharpe_ratio,drawdown,max_drawdown,yield_year_last,yield_year_first_to_last = \
                Analyze.evaluate_indicator(cur_nv,tradeday_per_year,anual_riskfree_rate) 
            indicators.loc[code,['年化收益率','年化波动率','夏普比率','最大回撤']] = [annual_yield,annual_std,sharpe_ratio,max_drawdown]
            
            drawdowns.append(drawdown) 
            yield_year_lasts.append(yield_year_last) 
            yield_year_first_to_lasts.append(yield_year_first_to_last) 
        
        drawdowns = pd.concat(drawdowns,axis=1); 
        yield_year_lasts = pd.concat(yield_year_lasts,axis=1)
        yield_year_first_to_lasts = pd.concat(yield_year_first_to_lasts,axis=1)
        
        return indicators,drawdowns,yield_year_lasts,yield_year_first_to_lasts

#回测
tradeday_per_year = 250

def backtest_stock(result):
    # 获取包含“收益”的列名
    columns = [col for col in result.columns if '收益' in col]
    # 计算平均收益率并添加新列
    result['平均收益率'] = result[columns].mean(axis=1)
    result['策略净值'] = (1+result['平均收益率']).cumprod()  #累乘
    result['策略净值'].iloc[0] = 1
    result['策略是否绝对获胜'] = (result['平均收益率']>0)
    absolute = result['策略是否绝对获胜'].sum()/result['策略是否绝对获胜'].count()
    _,_, sharpe_ratio, _, max_drawdown, _, _ = Analyze.evaluate_indicator(result['策略净值'], tradeday_per_year, 0)
    len_nav = result.shape[0]
    annual_yield = (result['策略净值'].iloc[-1]/result['策略净值'].iloc[0])**(tradeday_per_year/len_nav)-1
    return result, annual_yield, absolute, sharpe_ratio, max_drawdown

a=backtest_stock(result)
a[0].to_excel('回测结果.xlsx')

        
# 示例用法
'''
a = find_best_concept_stock('2016-01-06', '_884197_WI', index_members_dict, all_stock_dict, stock_num_select)

a=get_stock_date_pct('601727.SH', '2016-01-07', all_stock_dict)

sql = "SELECT S_INFO_WINDCODE, TRADE_DT, S_DQ_OPEN, S_DQ_CLOSE FROM ASHAREEODPRICES where S_INFO_WINDCODE='000509.SZ'"
stock = yunwind.read_sql(sql)

concept_stock = yunwind.read_sql("select * from AINDEXMEMBERSWIND where F_INFO_WINDCODE='884025.WI'")
concept_stock1 = yunwind.read_sql("select * from AINDEXDESCRIPTION where S_INFO_WINDCODE='884025.WI'")
'''
'''先求增长率，然后对应一下1的保留，nan的去掉，搞出一份dict
    每次找一下最大的记录列名，变成x*y列
def get_stock_pct(concept_index_sector):
    concept_index_members = pd.DataFrame(index=concept_index_sector.index, columns=concept_index_sector.columns)
    for code in concept_index_members.columns:
        sql = "select TRADE_DT, S_DQ_CLOSE from ASHAREEODPRICES where S_INFO_WINDCODE='{}'".format(code)
        stock = yunwind.read_sql(sql)
        stock['TRADE_DT'] = pd.to_datetime(stock['TRADE_DT'])  # 转换为日期类型
        stock.set_index('TRADE_DT', inplace=True)
        stock_pct = stock.pct_change()
        concept_index_members[code] = stock_pct.iloc[:, 0]
        print(code+'xxx')
    concept_index_members_filtered = concept_index_members.where(concept_index_sector == 1)
    return concept_index_members_filtered

new_dict = {}  # 创建一个新的字典
# 遍历原字典并处理 DataFrame
for key, df in concept_index_sector_dict.items():  # 使用 .items() 来遍历 key 和 df
    df_filter = get_stock_pct(df)  # 对 df 进行处理，得到 df_filter
    new_dict[key] = df_filter  # 将处理后的 df_filter 存入新字典中
    print(key)
'''

